
--[[
*********************************************************************************
** 
** Copyright 2021-2024 LuaOS
**
** THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
** IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
** FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
** AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
** LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
** OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
** SOFTWARE.
** 
*********************************************************************************
]]--

--------------------------------------------------------------------------------

local r_deliver         = rpc.r_deliver;
local r_bind            = rpc.r_bind;
local r_unbind          = rpc.r_unbind;
local r_response        = rpc.r_response;
local format            = string.format;
local insert            = table.insert;
local r_handlers        = {};
local r_readys          = {};
local cluster_nodes     = {};
local sockid_to_nodeid  = {};
local nodeid_to_sockid  = {};
local notify_delay      = {};
local notify_fname      = nil;
local notify_timer      = nil;
local node_identifier   = 0;
local const_max <const> = 0xffff;

local header = require("luaos.cluster.header");
local proto_type = require("luaos.cluster.protocol");

--------------------------------------------------------------------------------

local function encode(t)
  return proto_type.encode(t);
end

--------------------------------------------------------------------------------

local function decode(v)
  return proto_type.decode(v);
end

--------------------------------------------------------------------------------

local function is_local(who)
  return who <= const_max;
end

--------------------------------------------------------------------------------

local function set_ready(caller, info)
  if info then
    print(format("cluster node [%s] online!", info.module));
  elseif r_readys[caller] then
    print(format("cluster node [%s] leaves!", r_readys[caller].module));
  end
  r_readys[caller] = info;
end

--------------------------------------------------------------------------------

local function get_ready(caller)
  return r_readys[caller];
end

--------------------------------------------------------------------------------

local function rpc_update()
  for _, v in ipairs(notify_delay) do
    local n = rpc.dispatch(notify_fname, 0, 0, unwrap(v));
    if n == 0 then
      return 1000;
    end
  end
  notify_delay = {};
  notify_timer = nil;
end

--------------------------------------------------------------------------------

local function rpc_notify(ev, caller, ...)
  if not notify_fname then
    return;
  end
  if is_local(caller) then
    return;
  end
  if not notify_timer then
    local n = rpc.dispatch(notify_fname, 0, 0, ev, caller, ...);
    if n > 0 then
      return;
    end
    notify_timer = os.timer();
    notify_timer:expires(1000, rpc_update);
  end
  insert(notify_delay, wrap(ev, caller, ...));
end

--------------------------------------------------------------------------------

local function sendto_others(data)
  for id, session in pairs(cluster_nodes) do
    session.socket:send(data);
  end
end

--------------------------------------------------------------------------------

local function sendto_member(data, session)
  session.socket:send(data);
end

--------------------------------------------------------------------------------

local function lua_bind(info, caller, addr)
  local name = info.name;
  if not is_local(caller) then
    local from = info.module;
    r_bind(name, from, addr, caller, info.rcb);
  end
  if not r_handlers[caller] then
    r_handlers[caller] = {};
  end
  r_handlers[caller][name] = info;
  trace(format("bind %s by caller(%d)", name, caller));
end

--------------------------------------------------------------------------------

local function lua_unbind(name, caller)
  if not is_local(caller) then
    r_unbind(name, caller);
  end
  r_handlers[caller][name] = nil;
  trace(format("unbind %s by caller(%d)", name, caller));
end

--------------------------------------------------------------------------------

local function remove_session(peer)
  local id = peer:id();
  local session = cluster_nodes[id];
  if not session then
    return nil;
  end
  peer:close();
  local nodeid = sockid_to_nodeid[id];
  if nodeid then
    cluster_nodes [id] = nil;
    sockid_to_nodeid[id] = nil;
    nodeid_to_sockid[nodeid] = nil;
  end
  return nodeid;
end

--------------------------------------------------------------------------------

local function remove_handler(nodeid)
  local removed = {};
  for caller, v in pairs(r_handlers) do
    if not is_local(caller) then
      if caller & const_max == nodeid then
        removed[caller] = v;
      end
    end
  end
  for caller, v in pairs(removed) do
    local deleted = {};
    for name, _ in pairs(v) do
      insert(deleted, name);
    end
    for _, name in ipairs(deleted) do
      lua_unbind(name, caller);
    end
    r_handlers[caller] = nil;
  end
end

--------------------------------------------------------------------------------

local function remove_ready(nodeid)
  local removed = {};
  for caller, v in pairs(r_readys) do
    if not is_local(caller) then
      if caller & const_max == nodeid then
        removed[caller] = v;
      end
    end
  end
  for caller, v in pairs(removed) do
    set_ready(caller, nil);
    rpc_notify("leave", caller);
  end
end

--------------------------------------------------------------------------------

local function ws_on_error(peer, msg)
  local nodeid = remove_session(peer);
  if nodeid then
    remove_handler(nodeid);
    remove_ready(nodeid);
  end
end

--------------------------------------------------------------------------------

local function ws_on_receive(peer, data, ec)
  if ec then
    ws_on_error(peer, "receive error");
    return;
  end
  local id     = peer:id();
  local nodeid = sockid_to_nodeid[id];
  local addr   = cluster_nodes[id].ip;
  local info   = decode(data);
  local what   = info.what;
  
  if what == proto_type.deliver then
    local name   = info.name;
    local argv   = info.argv;
    local mask   = info.mask;
    local who    = info.who;
    local caller = info.caller << 16 | nodeid;
    local rcf    = info.rcf;
    local sn     = info.sn;
    r_deliver(name, argv, mask, who, caller, rcf, sn);
    return;
  end
  if what == proto_type.response then
    local data   = info.data;
    local caller = info.caller;
    local rcf    = info.rcf;
    local sn     = info.sn;
    r_response(data, caller, rcf, sn);
    return;
  end
  if what == proto_type.ready then
    local from   = info.module;
    local caller = info.caller << 16 | nodeid;
    if get_ready(caller) then
      return;
    end
    set_ready(caller, info);
    rpc_notify("online", caller, from, addr);
  end
  if what == proto_type.bind then
    local name   = info.name;
    local rcb    = info.rcb;
    local caller = info.caller << 16 | nodeid;    
    lua_bind(info, caller, addr);
    return;
  end
  if what == proto_type.unbind then
    local name   = info.name;
    local caller = info.caller << 16 | nodeid;
    lua_unbind(name, caller);
    return;
  end  
end

--------------------------------------------------------------------------------

local function lookout_handler(info)
  local what = info.what;  
  if what == proto_type.bind then
    lua_bind(info, info.caller);
    sendto_others(encode(info));
    return;
  end
  
  if what == proto_type.unbind then
    lua_unbind(info.name, info.caller);
    sendto_others(encode(info));
    return;
  end
  
  if what == proto_type.ready then
    sendto_others(encode(info));
    return;
  end
  
  local nodeid;
  if what == proto_type.deliver then
    nodeid = info.who & const_max;
    info.who = info.who >> 16;
  end
  
  if what == proto_type.response then
    nodeid = info.caller & const_max;
    info.caller = info.caller >> 16;
  end
  
  local id = nodeid_to_sockid[nodeid];
  local session = cluster_nodes[id];
  if session then
    sendto_member(encode(info), session);
  end
end

--------------------------------------------------------------------------------

local function sync_bounds(peer)
  for caller, bounds in pairs(r_handlers) do
    if is_local(caller) then
      for name, info in pairs(bounds) do
        peer:send(encode(info));
      end
      local ready = get_ready(caller);
      if ready then
        peer:send(encode(ready));
      end
    end
  end
end

--------------------------------------------------------------------------------

local function add_session(session, peerid)
  local id = session.socket:id();
  cluster_nodes[id] = session;
  sockid_to_nodeid[id] = peerid;
  nodeid_to_sockid[peerid] = id;
end

--------------------------------------------------------------------------------

local function new_session(peer, peerid)
  local what = peer:getheader(header.connect);
  if what ~= "luaos" then
    return false;
  end
  if not peerid then
    peerid = peer:getheader(header.identifier);
    peerid = tonumber(peerid);
  end
  local endpoint = peer:endpoint();
  local session = {
    socket = peer,
    peerid = peerid,
    ip     = endpoint.address,
    port   = endpoint.port,
  };
  add_session(session, peerid);
  sync_bounds(peer);
  peer:receive(bind(ws_on_receive, peer));
  return true;
end

--------------------------------------------------------------------------------

local function peer_exist(peerid)
  for _, session in pairs(cluster_nodes) do
    if session.peerid == peerid then
      return true;
    end
  end
  return false;
end

--------------------------------------------------------------------------------

local function ws_on_accept(peer, ec)
  if ec then
    ws_on_error(ec, peer, "accept error");
    return;
  end
  local peerid = nil;
  if not new_session(peer, peerid) then
    peer:close();
  end
end

--------------------------------------------------------------------------------

local function new_connect(host, port, selfid, peerid, protocol)
  if peer_exist(peerid) then
    return true;
  end
  local peer = io.socket(protocol);
  peer:setheader(header.connect, "luaos");
  peer:setheader(header.identifier, selfid);
  if peer:connect(host, port) then
    new_session(peer, peerid);
    return true;
  end
  return false;
end

--------------------------------------------------------------------------------

local function connect_members(socket, protocol)
  local selfid = 0;
  while true do
    local data = socket:read();
    if not data then
      return false;
    end
    local packet = decode(data);
    if packet.what == proto_type.number then
      selfid = packet.number;
    elseif packet.what == proto_type.ready then
      break;
    elseif packet.what ~= proto_type.forword then
      error(format("unknown protocol: %s", packet.what));
      return false;
    elseif selfid > 0 then
      local peerid = packet.id;
      local host   = packet.ip;
      local port   = packet.port;
      if not new_connect(host, port, selfid, peerid, protocol) then
        error(format("socket connect to %s:%d error", host, port));
        return false;
      end
      os.poll();
    end
  end
  print(format("cluster node number: %d", selfid));
  return true;
end

--------------------------------------------------------------------------------

local function listen_on_local(port, protocol)
  local server = io.server(protocol);
  local ok = server:listen("0.0.0.0", 0, ws_on_accept);
  if not ok then
    return;
  end
  return server, server:endpoint().port;
end

--------------------------------------------------------------------------------

local function on_connected(socket, context, ec)
  if ec then
    context.keepalive(socket, context, nil, ec);
    return;
  end
  socket:receive(bind(context.keepalive, socket, context));
  print("connection with cluster restored");
end

--------------------------------------------------------------------------------

local function on_reconnect(context)
  local socket = io.socket(context.protocol);
  socket:setheader(header.port, context.lport);
  socket:setheader(header.identifier, node_identifier);
  socket:connect(context.host, context.port, bind(on_connected, socket, context));
end

--------------------------------------------------------------------------------

local function keepalive(socket, context, data, ec)
  if ec then
    socket:close();
    local timer = os.timer();
    context.keepalive = keepalive;
    timer:expires(1000, bind(on_reconnect, context));
    error("connection error with cluster master");
  end
end

--------------------------------------------------------------------------------

local function connect_cluster(host, port, lport, alias, protocol)
  host = host or "127.0.0.1";
  port = port or 80;
  local socket = io.socket(protocol);
  node_identifier = alias or math.uuid();
  socket:setheader(header.port, lport);
  socket:setheader(header.identifier, node_identifier);
  if not socket:connect(host, port) then
    return nil;
  end
  local myinfo = socket:endpoint("local");
  if myinfo then
    storage.set("__nodeaddr__", myinfo.address);
  end
  rpc.lookout(lookout_handler);
  if not connect_members(socket, protocol) then
    return nil;
  end
  local context = {
    host     = host,
    port     = port,
    lport    = lport,
    protocol = protocol,
  };
  socket:receive(bind(keepalive, socket, context));
  return socket;
end

--------------------------------------------------------------------------------

local function main(host, port, rpc_name, alias)
  local protocol = "ws";
  local server, lport = listen_on_local(0, protocol);
  if not lport then
    error(format("socket listen error at port: %d", port));
    return;
  end  
  assert(rpc_name == nil or type(rpc_name) == "string");
  notify_fname = rpc_name;
  local socket = connect_cluster(host, port, lport, alias, protocol);
  if not socket then
    error(format("join to cluster error: %s:%d", host, port));
    return;
  end
  print(format("%s works on port %d", os.name(), lport));  
  while not os.stopped() do
    os.wait();
  end  
  rpc.lookout(nil);
  socket:close();
  server:close();
end

--------------------------------------------------------------------------------

luaos.start(main, ...);

--------------------------------------------------------------------------------
